import os

import openai
from langchain.chains import LLMChain, LLMMathChain
from langchain_community.llms import OpenAI
from langchain_core.prompts import PromptTemplate

def count_tokens(chain,query):
    result = chain.run(query=query)
    return result

os.environ["OPENAI_API_KEY"]=os.environ["OPENAI_API_KEY_ZHIHU"]
os.environ["OPENAI_API_BASE"]=os.environ["OPENAI_API_BASE_ZHIHU"]
llm = OpenAI(model_name= "gpt-3.5-turbo",temperature=0.9)

example_template = """{query}"""

example_prompt = PromptTemplate(input_variables=["query"],template=example_template)

llm_chain= LLMChain(llm=llm,prompt=example_prompt)
print(llm_chain(query="What is 13 raised to the 0.587965 power?"))

llm_math_chain = LLMMathChain(llm=llm,prompt=example_prompt)
count_tokens(llm_math_chain,query="What is 13 raised to the 0.587965 power")